import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from sklearn.compose import ColumnTransformer
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from imblearn.over_sampling import RandomOverSampler
from collections import Counter

# 修复画布错误
plt.switch_backend('TkAgg')  # 切换到TkAgg后端

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 数据加载
DATA_PATH = r"C:\Users\victus\Desktop\基本信息终版(4).xlsx"
data = pd.read_excel(DATA_PATH, usecols=['身高', '鞋码', '性别', '体重'])

# 数据预处理（增强过滤）
data = data.dropna()
data = data[(data['身高'] > 120) & (data['身高'] < 200)]
data = data[(data['鞋码'] > 32) & (data['鞋码'] < 48)]
data = data[(data['体重'] > 35) & (data['体重'] < 85)]  # 新增体重过滤

# 数据检查
print("数据总量:", len(data))
print("性别分布:", data['性别'].value_counts())

# 定义预处理管道
preprocessor = ColumnTransformer(
    transformers=[
        ('num', StandardScaler(), ['身高', '鞋码']),
        ('cat', OneHotEncoder(drop='first', handle_unknown='ignore'), ['性别'])
    ]
)

# 划分数据集
X = data[['身高', '鞋码', '性别']]
y = data['体重']
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=data['性别']
)

# 处理类别不平衡（回归问题使用样本权重）
ros = RandomOverSampler(random_state=42)
X_train_res, y_train_res = ros.fit_resample(X_train, y_train)

# 预处理流程
X_train_processed = preprocessor.fit_transform(X_train_res)
X_test_processed = preprocessor.transform(X_test)

# 模型训练（使用随机森林）
model = RandomForestRegressor(
    n_estimators=100,
    max_depth=5,
    min_samples_split=5,
    random_state=42
)
model.fit(X_train_processed, y_train_res)

# 模型评估
y_pred = model.predict(X_test_processed)

mse = mean_squared_error(y_test, y_pred)
mae = mean_absolute_error(y_test, y_pred)
r2 = r2_score(y_test, y_pred)

print("\n【模型性能】")
print(f"MSE: {mse:.2f} kg² | MAE: {mae:.2f} kg | R²: {r2:.2f}")

# 可视化（增强图表）
plt.figure(figsize=(12, 7))
plt.scatter(y_test, y_pred, alpha=0.7, c=X_test_processed[:, 0], cmap='viridis', edgecolors='k')
plt.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'k--', lw=2)
plt.xlabel('真实体重 (kg)')
plt.ylabel('预测体重 (kg)')
plt.title('随机森林体重预测')
plt.colorbar(label='身高标准化值')
plt.xticks(np.arange(40, 80, 5))
plt.yticks(np.arange(40, 80, 5))
plt.grid(True, linestyle='--', alpha=0.7)
plt.tight_layout()
plt.show()

# 模型保存
import joblib
joblib.dump(model, 'random_forest_model.pkl')
joblib.dump(preprocessor, 'preprocessor.pkl')